package com.learn.grpc.loadbalance;

import io.grpc.*;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static io.grpc.ConnectivityState.IDLE;

@Slf4j
public class LbLoadBalancer extends LoadBalancer {

    public static final Attributes.Key<LbRef<ConnectivityState>> STATE_INFO = Attributes.Key.create("state-info");

    private final Helper helper;

    // 保存当前地址连接
    private Map<EquivalentAddressGroup, Subchannel> subChannelMap = new ConcurrentHashMap<>();

    public LbLoadBalancer(Helper helper) {
        // 每一个 channel 会创建一个 LoadBalancer 对象
        this.helper = helper;
    }

    @Override
    public void handleResolvedAddresses(ResolvedAddresses resolvedAddresses) {
        // 服务实例变更时，调用 Listener2.onResult() 方法时会触发此方法执行

        // 解析服务最新的实例列表
        List<EquivalentAddressGroup> latestAddresses = resolvedAddresses
                .getAddresses()
                .stream()
                .flatMap(this::splitAddressCollection)
                .distinct()
                .collect(Collectors.toList());

        // 已经存在的地址
        Set<EquivalentAddressGroup> originAddresses = subChannelMap.keySet();

        // 1.对新的实例创建 SubChannel
        Map<EquivalentAddressGroup, Subchannel> newSubChannelMap = latestAddresses
                .stream()
                .filter(e -> !originAddresses.contains(e))
                .map(this::buildCreateSubChannelArgs)
                .map(helper::createSubchannel)
                .map(this::processSubChannel)
                .collect(Collectors.toConcurrentMap(Subchannel::getAddresses, s -> s));

        // 2.将已存在的 SubChannel 放到新的集合中
        originAddresses.stream()
                       .filter(latestAddresses::contains)
                       .forEach(e -> newSubChannelMap.put(e, subChannelMap.get(e)));

        // 3.将已下线的 SubChannel 关闭
        originAddresses.stream()
                       .filter(e -> !latestAddresses.contains(e))
                       .map(e -> subChannelMap.get(e))
                       .forEach(Subchannel::shutdown);

        // 替换为最新的 subChannelMap
        subChannelMap = newSubChannelMap;
    }

    private Stream<EquivalentAddressGroup> splitAddressCollection(EquivalentAddressGroup equivalentAddressGroup) {
        Attributes attributes = equivalentAddressGroup.getAttributes();
        return equivalentAddressGroup.getAddresses()
                .stream()
                .map(e -> new EquivalentAddressGroup(e, attributes));
    }


    private CreateSubchannelArgs buildCreateSubChannelArgs(EquivalentAddressGroup e) {
        return CreateSubchannelArgs.newBuilder()
                                   .setAddresses(e)
                                   .setAttributes(Attributes.newBuilder()
                                                            .set(STATE_INFO, new LbRef<>(IDLE))
                                                            .build())
                                   .build();
    }


    // 为每个 subChannel 创建状态监听器，并建立连接
    private Subchannel processSubChannel(Subchannel subchannel) {
        if (subChannelMap.containsValue(subchannel)) {
            return subchannel;
        }

        // 注册 subChannel 状态监听器
        subchannel.start(new LbSubChannelStateListener(this, subchannel, helper));
        // 建立连接
        subchannel.requestConnection();
        return subchannel;
    }


    @Override
    public void handleNameResolutionError(Status error) {
        helper.updateBalancingState(ConnectivityState.TRANSIENT_FAILURE, new LbSubChannelPicker(null));
    }

    @Override
    public void shutdown() {
        subChannelMap.values()
                     .stream()
                     .peek(s -> log.info("关闭" + s  + ", " + s.getAddresses()))
                     .forEach(Subchannel::shutdown);
    }

    public Map<EquivalentAddressGroup, Subchannel> getSubChannelMap() {
        return new ConcurrentHashMap<>(this.subChannelMap);
    }
}
